package com.joker.demo.controller;


import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.grpc.Collections;
import jakarta.annotation.Resource;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;

/**
 * 向量数据库测试
 *
 * @author : feixiang.li
 * @since : 2025-09-25 21:39
 */
@RestController
@RequestMapping("/embedding")
public class EmbeddingController {


    @Resource
    EmbeddingModel embeddingModel;


    @RequestMapping("/embedding")
    public String embedding() {
        Response<Embedding> embeddingResponse = embeddingModel.embed("hello world");
        System.out.println(embeddingResponse);
        return embeddingResponse.content().vectorAsList().toString();
    }


    @Resource
    private QdrantClient qdrantClient;

    /**
     * 创建向量数据库
     *
     * @return
     */
    @RequestMapping("/createQdrant")
    public String createQdrant() {
        var build = Collections.VectorParams.newBuilder()
                .setDistance(Collections.Distance.Cosine)
                .setSize(1024)
                .build();
        qdrantClient.createCollectionAsync("my-qiniu", build);
        return "创建成功";
    }


    @Resource
    EmbeddingStore<TextSegment> embeddingStore;

    /**
     * 新增记录
     */
    @RequestMapping("/addRecord")
    public String addRecord() {
        TextSegment textSegment = TextSegment.from("飞哥写了一首诗，太阳当空照，我要找工作，鸡鸣破晓光，");
        textSegment.metadata().put("id", "1");
        Embedding embedding = embeddingModel.embed(textSegment).content();
        String result = embeddingStore.add(embedding, textSegment);

        System.out.println(result);

        textSegment = TextSegment.from("我叫斗破苍穹，现在等级划分，练气期，金丹境，大乘境");
        textSegment.metadata().put("id", "2");
        embedding = embeddingModel.embed(textSegment).content();
        result = embeddingStore.add(embedding, textSegment);

        System.out.println(result);
        return "新增成功";
    }


    @RequestMapping("/search")
    public String search() {
        Embedding queryEmbedding = embeddingModel.embed("斗破的境界都有啥？").content();
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .maxResults(1)
                .build();
        EmbeddingSearchResult<TextSegment> search = embeddingStore.search(embeddingSearchRequest);
        System.out.println(search.matches().get(0).embedded().text());
        return "查询成功";
    }

    @RequestMapping("/search2")
    public String search2() {
        Embedding queryEmbedding = embeddingModel.embed("斗破的境界都有啥？").content();
        EmbeddingSearchRequest embeddingSearchRequest = EmbeddingSearchRequest.builder()
                .queryEmbedding(queryEmbedding)
                .filter(metadataKey("id").isEqualTo("2"))
                .maxResults(1)
                .build();
        EmbeddingSearchResult<TextSegment> search = embeddingStore.search(embeddingSearchRequest);
        System.out.println(search.matches().get(0).embedded().text());
        return "查询成功";
    }
}
